
#include "ECDSASigning.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block signing and not pipe signing.
 * 
 * @param ecdsaprvkey The ECDSA private key to use to sign. */
ECDSASigning::ECDSASigning(const ECDSAPrivateKey &ecdsaprvkey)
{
   ecdsaprvkey.getSSLKeyComponents(&key);
   input = NULL;
   output = NULL;
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe signing.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param ecdsaprvkey The ECDSA private key to use to sign the
 * data. */
ECDSASigning::ECDSASigning(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const ECDSAPrivateKey &ecdsaprvkey)
{
   ecdsaprvkey.getSSLKeyComponents(&key);
   input = &iinput;
   output = &ioutput;
}

/** Destructor for the ECDSA private key signing class. */
ECDSASigning::~ECDSASigning()
{
   EC_KEY_free(key);
}

/** Sign data block and return signed data. The signed data returned
 * must be freed by the caller.
 * 
 * @param datablock A block of data to sign.
 * 
 * @param datasize The length of the data block to sign.
 * 
 * @param[out] signedsize The length of the signed block of data
 * returned. The data returned must be freed by the caller. */
char *ECDSASigning::processDataBlock(const char *datablock, apr_size_t datasize, apr_size_t *signedsize)
{
   if (datablock == NULL)
   {
      errprint("tried to sign NULL data block");
      return NULL;
   }
   if (key == NULL)
   {
      errprint("key is NULL - signing failed");
      return NULL;
   }
   EVP_MD_CTX ctxt;
   EVP_MD_CTX_init(&ctxt);
   if (!EVP_SignInit_ex(&ctxt, EVP_sha256(), NULL))
   {
      errprint("could not initialise signing");
      return NULL;
   }
   if (!EVP_SignUpdate(&ctxt, (void *)datablock, datasize))
   {
      errprint("could not add data to signing");
      return NULL;
   }
   EVP_PKEY *newpkey = EVP_PKEY_new();
   if (newpkey == NULL)
   {
      errprint("could not initialise the private key");
      return NULL;
   }
   if (key)
      EVP_PKEY_set1_EC_KEY(newpkey, key);
   unsigned int sign_length = EVP_PKEY_size(newpkey);
   char *signature = (char *)malloc(sizeof(char)*sign_length);
   if (!EVP_SignFinal(&ctxt, (unsigned char *)signature, &sign_length, newpkey))
   {
      errprint("could not finalise signing process");
      EVP_PKEY_free(newpkey);
      free(signature);
      return NULL;
   }
   char *final_data = (char *)malloc(sizeof(char)*(datasize + sign_length + sizeof(unsigned int)));
   memcpy(final_data, datablock, sizeof(char)*datasize);
   memcpy(&(final_data[datasize]), signature, sizeof(char)*sign_length);
   unsigned int network_siglength =
      ((uint32_t)(((char *)&sign_length)[0]) << 24) |
      ((uint32_t)(((char *)&sign_length)[1]) << 16) |
      ((uint32_t)(((char *)&sign_length)[2]) << 8) |
      ((uint32_t)(((char *)&sign_length)[3]));
   memcpy(&(final_data[datasize + sign_length]), &network_siglength, sizeof(unsigned int));
   if (signedsize)
      *signedsize = datasize + sign_length + sizeof(unsigned int);
   EVP_PKEY_free(newpkey);
   EVP_MD_CTX_cleanup(&ctxt);
   free(signature);
   return final_data;
}

/** Begin the process of signing a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int ECDSASigning::processPipeBegin()
{
   EVP_MD_CTX_init(&e_ctxt);
   if (!EVP_SignInit_ex(&e_ctxt, EVP_sha256(), NULL))
   {
      errprint("could not initialise signing");
      return 1;
   }
   return 0;
}

/** Process any outstanding plaintext data on the input pipe, and
 * place any processed plaintext data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int ECDSASigning::processPipeStep()
{
   int input_block = input->realLength();
   char *itemp_block = (char *)malloc(sizeof(char)*input_block);

   apr_size_t readbytes = input->readFromStream(itemp_block, input_block);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != input_block)
   {
      errprint("failed to read from stream (read %" APR_SIZE_T_FMT ", expected %" APR_SIZE_T_FMT ")", readbytes, (apr_size_t)input_block);
      free(itemp_block);
      return 1;
   }        
   if (!EVP_SignUpdate(&e_ctxt, (void *)itemp_block, input_block))
   {
      errprint("could not add data to signing");
      free(itemp_block);
      return 1;
   }
   apr_size_t writebytes = output->writeToStream(itemp_block, input_block);
   if (writebytes != input_block)
   {
      errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(input_block));
      free(itemp_block);
      return 1;
   }
   input->removeFromStream(readbytes);
   free(itemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int ECDSASigning::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding plaintext data in the pipe, process any
 * outstanding data in the internal buffers of the signing process and
 * append the signature to the stream.
 * 
 * @return Zero on success, non-zero on failure. */
int ECDSASigning::processPipeEnd()
{
   processPipeStep();
   EVP_PKEY *newpkey = EVP_PKEY_new();
   if (newpkey == NULL)
   {
      errprint("could not initialise the private key");
      EVP_PKEY_free(newpkey);
      return 1;
   }
   if (key)
      EVP_PKEY_set1_EC_KEY(newpkey, key);
   unsigned int sign_length = EVP_PKEY_size(newpkey);
   char *signature = (char *)malloc(sizeof(char)*sign_length);
   if (!EVP_SignFinal(&e_ctxt, (unsigned char *)signature, &sign_length, newpkey))
   {
      errprint("could not finalise signing process");
      EVP_PKEY_free(newpkey);
      EVP_MD_CTX_cleanup(&e_ctxt);
      free(signature);
      return 1;
   }
   output->writeToStream(signature, sign_length);
   unsigned int network_siglength =
      ((uint32_t)(((char *)&sign_length)[0]) << 24) |
      ((uint32_t)(((char *)&sign_length)[1]) << 16) |
      ((uint32_t)(((char *)&sign_length)[2]) << 8) |
      ((uint32_t)(((char *)&sign_length)[3]));
   output->writeToStream((char *)&network_siglength, sizeof(unsigned int));
   EVP_PKEY_free(newpkey);
   EVP_MD_CTX_cleanup(&e_ctxt);
   free(signature);
   return 0;
}
